package scu.maqiang.heat;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES2T31;
import scu.maqiang.mesh.Mesh2T3;
import scu.maqiang.mesh.ScalarFunc;
import scu.maqiang.mesh.T3Type;
import scu.maqiang.numeric.Direct;
import scu.maqiang.numeric.DirectSSolver;
import scu.maqiang.numeric.NewIterSSolver;
import scu.maqiang.numeric.SRMatrix;

/**
 * @brief 二维热传导问题通用计算类
 * @author 马强
 *
 */
public class Heat2DT31 {
	public Heat2DT31(Mesh2T3 mesh) {
		this.mesh = mesh;
		fs = new FES2T31(mesh);
		A = new SRMatrix(fs.GetNdof());
		M = new SRMatrix(fs.GetNdof());
		RHS = new double[fs.GetNdof()];
		source = new double[fs.GetNdof()];
		flux1 = new double[fs.GetNdof()];
		flux2 = new double[fs.GetNdof()];
		gradient = new double[2][fs.GetNdof()];
	}

	public static void main(String[] args) {
		long begin = System.nanoTime();
		double theta = 0.5;
		double T = 2.0;
		double dt = T / 100.0;
		Mesh2T3 mesh = new Mesh2T3().square2D(10, 10, T3Type.Left);

		Heat2DT31 heat = new Heat2DT31(mesh);
		ScalarFunc initFunc = (xy, label, param) -> 0.0;
		double[] x0 = mesh.functionInMesh(initFunc, null);
		heat.GeneralHeatProblem(new BVPType[] { BVPType.COMMON, BVPType.CONSIST_MASS}, 1.0, 1.0, 1.0,
				new double[] { theta, T, dt }, new int[] { 1, 2, 3, 4 }, 0.0, null, 0.0, null, 0.0, 0.0, x0, 
				"MyHeat2DT31.dat");
		long end = System.nanoTime();
		System.out.println("Solve time: " + (end - begin) / 1.0e9 + " s");
	}

	/**
	 * @brief 一般瞬态热传导计算程序, 问题中所有的参数均为常数
	 * @param tp                     问题类型, 二维向量, 可选择组合为 
	 *                               {COMMON, CONSIST_MATRIX/LUMP_MATRIX}:
	 * 								 {AXISYMMETRIC, AXIS_CONSIST_MATRIX/AXIS_LUMP_MATRIX}:
	 * @param kxx                    材料热传导系数
	 * @param capacity               材料热容系数， 为材料密度与比热乘积
	 * @param sourceValue            热源值
	 * @param timeIntegration        时间积分系数, 使用Theta格式, 三维向量: {Theta, T ,dt}, 其中
	 *                               0<=Theta<=1, Theta=0时为全显示格式, Theta=1为全隐式格式, Theta=0.5表示Crank-Nicolson格式，
	 *                               T为终止时间, dt为时间步长
	 * @param DirichletBoundaryLabel Dirichlet边界编号
	 * @param fixedvalue             Dirichlet边界温度值
	 * @param NeumannBoundaryLabel   热流边界编号
	 * @param fluxValue              热流边界值
	 * @param RobinBoundaryLabel     Robin边界编号
	 * @param robinFluxValue         Robin边界热流值
	 * @param convectValue           Robin边界对流换热系数
	 * @param initValue				  初始值向量
	 * @param reportFile             最终报告文件
	 */
	public void GeneralHeatProblem(BVPType[] tp, double kxx, double capacity, double sourceValue,
			double[] timeIntegration, int[] DirichletBoundaryLabel, double fixedvalue, int[] NeumannBoundaryLabel,
			double fluxValue, int[] RobinBoundaryLabel, double robinFluxValue, double convectValue, 
			double[] initValue, String reportFile) {
		double theta = timeIntegration[0];
		double T = timeIntegration[1];
		double dt = timeIntegration[2];
		fs.assembleStiff(new double[]{kxx}, tp[0], A);
		if (RobinBoundaryLabel != null) {
			fs.assembleBoundaryMass(new double[] { convectValue }, tp[1], A, RobinBoundaryLabel);
		}
		fs.assembleMass(new double[]{capacity}, tp[1], M);

		SRMatrix L = M.add(theta * dt, A);
		SRMatrix R = M.add((theta - 1) * dt, A);
		if (Math.abs(sourceValue) > 0.0) {
			fs.assembleSource(new double[]{sourceValue}, tp[0], source);
		}
		if (NeumannBoundaryLabel != null) {
			fs.assembleFlux(new double[]{fluxValue}, tp[0], flux1, NeumannBoundaryLabel);
		}
		if (RobinBoundaryLabel != null) {
			fs.assembleFlux(new double[]{robinFluxValue}, tp[0], flux2, RobinBoundaryLabel);
		}

		boolean dirichletExist = DirichletBoundaryLabel != null;
		if (dirichletExist) {
			fs.applyBC_MBN(L, Direct.All, DirichletBoundaryLabel);
		}

		//NewIterSSolver solver = new NewIterSSolver(L);
		// DirectSSolver solver = new DirectSSolver(L, fs.getnBand());
		DirectSSolver solver = new DirectSSolver(L);
		solver.Factorize_Cholesky(1);
		double time = 0.0;
		x = initValue;
		mesh.toTecplot(reportFile, time, x);
		double[] Rx = new double[fs.GetNdof()];
		while (time < T) {
			time += dt;
			System.out.println("time = " + time);
			R.mul(x, Rx);
			for (int i = 0, n = fs.GetNdof(); i < n; i++) {
				RHS[i] = Rx[i] + (flux1[i] + flux2[i] + source[i]) * dt;
			}
			if (dirichletExist) {
				fs.applyBC_MBN(RHS, Direct.All, fixedvalue, DirichletBoundaryLabel);
			}
			//solver.PCGSSOR(RHS, x, 1.50, 1);
			x = solver.Solve_Cholesky(RHS, 1);
			// solver.PLUSolve(RHS, x, 1);
			mesh.toTecplot(reportFile, time, x);
		}
	}

	public void GeneralHeatProblem(BVPType[] tp, ScalarFunc kxxFunc, double[] paramKxx, ScalarFunc capacityFunc,
			double[] paramCapacity, ScalarFunc sourceFunc, double[] paramSource, double[] timeIntegration,
			int[] DirichletBoundaryLabel, ScalarFunc fixedFunc, double[] paramFixed,
								   int[] NeumannBoundaryLabel,
			ScalarFunc fluxFunc, double[] paramFlux, int[] RobinBoundaryLabel, ScalarFunc robinFluxFunc,
			double[] paramRobinFlux, ScalarFunc convectFunc, double[] paramConvectFunc, 
			double[] initFunc, String reportFile) {
		double theta = timeIntegration[0];
		double T = timeIntegration[1];
		double dt = timeIntegration[2];

		fs.assembleStiff(new ScalarFunc[] { kxxFunc }, paramKxx, tp[0], A);
		if (RobinBoundaryLabel != null) {
			fs.assembleBoundaryMass(new ScalarFunc[] { convectFunc }, paramConvectFunc, tp[1], A,
					RobinBoundaryLabel);
		}
		fs.assembleMass(new ScalarFunc[] { capacityFunc }, paramCapacity, tp[1], M);

		SRMatrix L = M.add(theta * dt, A);
		SRMatrix R = M.add((theta - 1) * dt, A);

		paramSource[0] = 0.0;
		boolean sourceExist = sourceFunc != null;
		if (sourceExist) {
			fs.assembleSource(new ScalarFunc[] { sourceFunc }, paramSource, tp[0], source);
		}

		double[] source2 = new double[fs.GetNdof()];

		boolean NeumannExist = NeumannBoundaryLabel != null;
		if (NeumannExist) {
			paramFlux[0] = 0.0;
			fs.assembleFlux(new ScalarFunc[]{fluxFunc}, paramFlux, tp[0], flux1, NeumannBoundaryLabel);
		}

		double[] flux11 = new double[fs.GetNdof()];

		boolean RobinExist = RobinBoundaryLabel != null;
		if (RobinExist) {
			paramRobinFlux[0] = 0.0;
			fs.assembleFlux(new ScalarFunc[]{robinFluxFunc}, paramRobinFlux, tp[0], flux2, RobinBoundaryLabel);
		}

		double[] flux22 = new double[fs.GetNdof()];

		boolean dirichletExist = DirichletBoundaryLabel != null;
		if (dirichletExist) {
			fs.applyBC_MBN(L, Direct.All, DirichletBoundaryLabel);
		}

		NewIterSSolver solver = new NewIterSSolver(L);
		double time = 0.0;
		x = initFunc;
		mesh.toTecplot(reportFile, time, x);
		double[] Rx = new double[fs.GetNdof()];
		while (time < T) {
			time += dt;
			System.out.println("time = " + time);
			R.mul(x, Rx);
			if (sourceExist) {
				paramSource[0] = time;
				fs.assembleSource(new ScalarFunc[] { sourceFunc }, paramSource, tp[0], source2);
			}

			if (NeumannExist) {
				paramFlux[0] = time;
				fs.assembleFlux(new ScalarFunc[]{fluxFunc}, paramFlux, tp[0], flux11, NeumannBoundaryLabel);
			}

			if (RobinExist) {
				paramRobinFlux[0] = time;
				fs.assembleFlux(new ScalarFunc[]{robinFluxFunc}, paramRobinFlux, tp[0], flux22, RobinBoundaryLabel);
			}

			for (int i = 0, n = fs.GetNdof(); i < n; i++) {
				RHS[i] = Rx[i] + (1 - theta) * (flux11[i] + flux22[i] + source2[i]) * dt
						+ theta * (flux1[i] + flux2[i] + source[i]) * dt;
				flux1[i] = flux11[i];
				flux2[i] = flux22[i];
				source[i] = source2[i];
			}
			if (dirichletExist) {
				paramFixed[0] = time;
				fs.applyBC_MBN(RHS, Direct.All, new ScalarFunc[] { fixedFunc }, paramFixed, DirichletBoundaryLabel);
			}
			solver.PCGSSOR(RHS, x, 1.50, 1);
			mesh.toTecplot(reportFile, time, x);
		}
	}

	Mesh2T3 mesh;
	FES2T31 fs;
	SRMatrix A;
	SRMatrix M;
	double[] source;
	double[] flux1;
	double[] flux2;
	double[] RHS;
	double[] x;
	double[][] gradient;
}
